import torch
from torch import nn
import torch.nn.functional as F

__all__ = ['UNet', 'NestedUNet', 'EnhancedNestedUNet', 'EnhancedNestedUNetDualDomain']


class VGGBlock(nn.Module):
    def __init__(self, in_channels, middle_channels, out_channels):
        super().__init__()
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_channels, middle_channels, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(middle_channels)
        self.conv2 = nn.Conv2d(middle_channels, out_channels, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        return out


class SEBlock(nn.Module):
    def __init__(self, channels, reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        reduced = max(channels // reduction, 1)
        self.fc = nn.Sequential(
            nn.Linear(channels, reduced, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(reduced, channels, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y


class sSEBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, 1, kernel_size=1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        attn = self.sigmoid(self.conv(x))
        return x * attn


class scSEBlock(nn.Module):
    def __init__(self, channels, reduction=16):
        super().__init__()
        self.cSE = SEBlock(channels, reduction)
        self.sSE = sSEBlock(channels)

    def forward(self, x):
        return self.cSE(x) + self.sSE(x)


class ResSEBlock(nn.Module):
    def __init__(self, in_channels, middle_channels, out_channels, reduction=16):
        super().__init__()
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_channels, middle_channels, 3, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(middle_channels)
        self.conv2 = nn.Conv2d(middle_channels, out_channels, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.scse = scSEBlock(out_channels, reduction)
        self.downsample = None
        if in_channels != out_channels:
            self.downsample = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        out = self.scse(out)

        if self.downsample is not None:
            identity = self.downsample(identity)

        out = self.relu(out + identity)
        return out


class DualDomainFusion(nn.Module):
    def __init__(self, channels):
        super().__init__()
        reduced = max(channels // 4, 8)
        self.spatial = nn.Sequential(
            nn.Conv2d(channels, channels, kernel_size=3, padding=1, groups=channels, bias=False),
            nn.BatchNorm2d(channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels, channels, kernel_size=1, bias=False),
            nn.Sigmoid()
        )
        self.freq = nn.Sequential(
            nn.Conv2d(channels, reduced, kernel_size=1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(reduced, channels, kernel_size=1, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        spatial_weight = self.spatial(x)
        freq = torch.fft.fft2(x, norm='ortho')
        freq_mag = freq.abs()
        freq_weight = self.freq(freq_mag)
        weight = 0.5 * spatial_weight + 0.5 * freq_weight
        return x + x * weight


class CrossScaleGatedFusion(nn.Module):
    def __init__(self, in_channels, branches):
        super().__init__()
        hidden = max(in_channels // 4, 8)
        self.branches = branches
        self.gate = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, hidden, kernel_size=1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(hidden, branches, kernel_size=1, bias=True)
        )

    def forward(self, features):
        concat = torch.cat(features, dim=1)
        weights = self.gate(concat)
        weights = torch.softmax(weights.view(weights.size(0), self.branches, 1, 1), dim=1)
        weighted = []
        for idx, feat in enumerate(features):
            weighted.append(feat * weights[:, idx:idx + 1])
        return torch.cat(weighted, dim=1)


class HybridContextModule(nn.Module):
    def __init__(self, channels):
        super().__init__()
        branch_ch = max(channels // 2, 8)
        self.branch1 = nn.Sequential(
            nn.Conv2d(channels, branch_ch, kernel_size=1, bias=False),
            nn.BatchNorm2d(branch_ch),
            nn.ReLU(inplace=True),
        )
        self.branch2 = nn.Sequential(
            nn.Conv2d(channels, branch_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(branch_ch),
            nn.ReLU(inplace=True),
        )
        self.branch3 = nn.Sequential(
            nn.Conv2d(channels, branch_ch, kernel_size=3, padding=3, dilation=3, bias=False),
            nn.BatchNorm2d(branch_ch),
            nn.ReLU(inplace=True),
        )
        self.global_pool = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(channels, branch_ch, kernel_size=1, bias=False),
            nn.BatchNorm2d(branch_ch),
            nn.ReLU(inplace=True),
        )
        self.mix = nn.Sequential(
            nn.Conv2d(branch_ch * 4, channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(channels),
            nn.ReLU(inplace=True)
        )
        self.dual = DualDomainFusion(channels)

    def forward(self, x):
        g = self.global_pool(x)
        g = F.interpolate(g, size=x.shape[2:], mode='bilinear', align_corners=True)
        out = torch.cat([self.branch1(x), self.branch2(x), self.branch3(x), g], dim=1)
        out = self.mix(out)
        return self.dual(out)


class HarmonicResSEBlock(nn.Module):
    def __init__(self, in_channels, out_channels, reduction=16, dilation=1):
        super().__init__()
        self.base = ResSEBlock(in_channels, out_channels, out_channels, reduction)
        self.dw = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=dilation,
                            dilation=dilation, groups=out_channels, bias=False)
        self.dw_bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.dual = DualDomainFusion(out_channels)

    def forward(self, x):
        out = self.base(x)
        dw = self.dw(out)
        dw = self.dw_bn(dw)
        out = self.relu(out + dw)
        return self.dual(out)


class FusionHarmonicBlock(nn.Module):
    def __init__(self, branch_channels, out_channels, reduction=16, dilation=1):
        super().__init__()
        total_in = sum(branch_channels)
        self.fusion = CrossScaleGatedFusion(total_in, len(branch_channels))
        self.block = HarmonicResSEBlock(total_in, out_channels, reduction, dilation)

    def forward(self, features):
        fused = self.fusion(features)
        return self.block(fused)


class UNet(nn.Module):
    def __init__(self, num_classes, input_channels=3, **kwargs):
        super().__init__()

        nb_filter = [32, 64, 128, 256, 512]

        self.pool = nn.MaxPool2d(2, 2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
        self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1])
        self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2])
        self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3])
        self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4])

        self.conv3_1 = VGGBlock(nb_filter[3]+nb_filter[4], nb_filter[3], nb_filter[3])
        self.conv2_2 = VGGBlock(nb_filter[2]+nb_filter[3], nb_filter[2], nb_filter[2])
        self.conv1_3 = VGGBlock(nb_filter[1]+nb_filter[2], nb_filter[1], nb_filter[1])
        self.conv0_4 = VGGBlock(nb_filter[0]+nb_filter[1], nb_filter[0], nb_filter[0])

        self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)


    def forward(self, input):
        x0_0 = self.conv0_0(input)
        x1_0 = self.conv1_0(self.pool(x0_0))
        x2_0 = self.conv2_0(self.pool(x1_0))
        x3_0 = self.conv3_0(self.pool(x2_0))
        x4_0 = self.conv4_0(self.pool(x3_0))

        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
        x2_2 = self.conv2_2(torch.cat([x2_0, self.up(x3_1)], 1))
        x1_3 = self.conv1_3(torch.cat([x1_0, self.up(x2_2)], 1))
        x0_4 = self.conv0_4(torch.cat([x0_0, self.up(x1_3)], 1))

        output = self.final(x0_4)
        return output


class NestedUNet(nn.Module):
    def __init__(self, num_classes, input_channels=3, deep_supervision=False, **kwargs):
        super().__init__()

        nb_filter = [32, 64, 128, 256, 512]

        self.deep_supervision = deep_supervision

        self.pool = nn.MaxPool2d(2, 2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
        self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1])
        self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2])
        self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3])
        self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4])

        self.conv0_1 = VGGBlock(nb_filter[0]+nb_filter[1], nb_filter[0], nb_filter[0])
        self.conv1_1 = VGGBlock(nb_filter[1]+nb_filter[2], nb_filter[1], nb_filter[1])
        self.conv2_1 = VGGBlock(nb_filter[2]+nb_filter[3], nb_filter[2], nb_filter[2])
        self.conv3_1 = VGGBlock(nb_filter[3]+nb_filter[4], nb_filter[3], nb_filter[3])

        self.conv0_2 = VGGBlock(nb_filter[0]*2+nb_filter[1], nb_filter[0], nb_filter[0])
        self.conv1_2 = VGGBlock(nb_filter[1]*2+nb_filter[2], nb_filter[1], nb_filter[1])
        self.conv2_2 = VGGBlock(nb_filter[2]*2+nb_filter[3], nb_filter[2], nb_filter[2])

        self.conv0_3 = VGGBlock(nb_filter[0]*3+nb_filter[1], nb_filter[0], nb_filter[0])
        self.conv1_3 = VGGBlock(nb_filter[1]*3+nb_filter[2], nb_filter[1], nb_filter[1])

        self.conv0_4 = VGGBlock(nb_filter[0]*4+nb_filter[1], nb_filter[0], nb_filter[0])

        if self.deep_supervision:
            self.final1 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final2 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final3 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final4 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
        else:
            self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)


    def forward(self, input):
        x0_0 = self.conv0_0(input)
        x1_0 = self.conv1_0(self.pool(x0_0))
        x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_0)], 1))

        x2_0 = self.conv2_0(self.pool(x1_0))
        x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_0)], 1))
        x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.up(x1_1)], 1))

        x3_0 = self.conv3_0(self.pool(x2_0))
        x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_0)], 1))
        x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.up(x2_1)], 1))
        x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.up(x1_2)], 1))

        x4_0 = self.conv4_0(self.pool(x3_0))
        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
        x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1))
        x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1))
        x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 1))

        if self.deep_supervision:
            output1 = self.final1(x0_1)
            output2 = self.final2(x0_2)
            output3 = self.final3(x0_3)
            output4 = self.final4(x0_4)
            return [output1, output2, output3, output4]

        else:
            output = self.final(x0_4)
            return output


class EnhancedNestedUNet(nn.Module):
    def __init__(self, num_classes, input_channels=3, deep_supervision=False, width_multiplier=1.0, **kwargs):
        super().__init__()

        base = [32, 64, 128, 256, 512]
        nb_filter = [int(c * width_multiplier) for c in base]

        self.deep_supervision = deep_supervision

        self.pool = nn.MaxPool2d(2, 2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        # Encoder with residual scSE blocks
        self.conv0_0 = ResSEBlock(input_channels, nb_filter[0], nb_filter[0])
        self.conv1_0 = ResSEBlock(nb_filter[0], nb_filter[1], nb_filter[1])
        self.conv2_0 = ResSEBlock(nb_filter[1], nb_filter[2], nb_filter[2])
        self.conv3_0 = ResSEBlock(nb_filter[2], nb_filter[3], nb_filter[3])
        self.conv4_0 = ResSEBlock(nb_filter[3], nb_filter[4], nb_filter[4])

        # ASPP at bottleneck to enhance multi-scale features
        self.aspp = nn.Sequential(
            nn.Conv2d(nb_filter[4], nb_filter[4] // 4, kernel_size=1, bias=False),
            nn.BatchNorm2d(nb_filter[4] // 4),
            nn.ReLU(inplace=True),
            nn.Conv2d(nb_filter[4] // 4, nb_filter[4] // 4, kernel_size=3, padding=2, dilation=2, bias=False),
            nn.BatchNorm2d(nb_filter[4] // 4),
            nn.ReLU(inplace=True),
            nn.Conv2d(nb_filter[4] // 4, nb_filter[4], kernel_size=1, bias=False),
            nn.BatchNorm2d(nb_filter[4]),
            nn.ReLU(inplace=True)
        )

        # Decoder dense skip with residual scSE blocks (UNet++ style)
        self.conv0_1 = ResSEBlock(nb_filter[0] + nb_filter[1], nb_filter[0], nb_filter[0])
        self.conv1_1 = ResSEBlock(nb_filter[1] + nb_filter[2], nb_filter[1], nb_filter[1])
        self.conv2_1 = ResSEBlock(nb_filter[2] + nb_filter[3], nb_filter[2], nb_filter[2])
        self.conv3_1 = ResSEBlock(nb_filter[3] + nb_filter[4], nb_filter[3], nb_filter[3])

        self.conv0_2 = ResSEBlock(nb_filter[0] * 2 + nb_filter[1], nb_filter[0], nb_filter[0])
        self.conv1_2 = ResSEBlock(nb_filter[1] * 2 + nb_filter[2], nb_filter[1], nb_filter[1])
        self.conv2_2 = ResSEBlock(nb_filter[2] * 2 + nb_filter[3], nb_filter[2], nb_filter[2])

        self.conv0_3 = ResSEBlock(nb_filter[0] * 3 + nb_filter[1], nb_filter[0], nb_filter[0])
        self.conv1_3 = ResSEBlock(nb_filter[1] * 3 + nb_filter[2], nb_filter[1], nb_filter[1])

        self.conv0_4 = ResSEBlock(nb_filter[0] * 4 + nb_filter[1], nb_filter[0], nb_filter[0])

        self.dropout2d = nn.Dropout2d(p=0.2)

        if self.deep_supervision:
            self.final1 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final2 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final3 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final4 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
        else:
            self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)

    def forward(self, input):
        # Encoder
        x0_0 = self.conv0_0(input)
        x1_0 = self.conv1_0(self.pool(x0_0))
        x2_0 = self.conv2_0(self.pool(x1_0))
        x3_0 = self.conv3_0(self.pool(x2_0))
        x4_0 = self.conv4_0(self.pool(x3_0))
        x4_0 = self.aspp(x4_0)

        # Decoder UNet++ dense connections
        x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_0)], 1))
        x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_0)], 1))
        x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_0)], 1))
        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))

        x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.up(x1_1)], 1))
        x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.up(x2_1)], 1))
        x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1))

        x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.up(x1_2)], 1))
        x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1))

        x0_3 = self.dropout2d(x0_3)
        x1_3 = self.dropout2d(x1_3)

        x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 1))

        if self.deep_supervision:
            output1 = self.final1(x0_1)
            output2 = self.final2(x0_2)
            output3 = self.final3(x0_3)
            output4 = self.final4(x0_4)
            return [output1, output2, output3, output4]
        else:
            output = self.final(x0_4)
            return output


class EnhancedNestedUNetDualDomain(nn.Module):
    def __init__(self, num_classes, input_channels=3, deep_supervision=False, width_multiplier=1.0, **kwargs):
        super().__init__()

        base = [32, 64, 128, 256, 512]
        nb_filter = [int(c * width_multiplier) for c in base]

        self.deep_supervision = deep_supervision

        self.pool = nn.MaxPool2d(2, 2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        # Encoder with harmonic residual attention blocks
        self.conv0_0 = HarmonicResSEBlock(input_channels, nb_filter[0], dilation=1)
        self.conv1_0 = HarmonicResSEBlock(nb_filter[0], nb_filter[1], dilation=1)
        self.conv2_0 = HarmonicResSEBlock(nb_filter[1], nb_filter[2], dilation=2)
        self.conv3_0 = HarmonicResSEBlock(nb_filter[2], nb_filter[3], dilation=2)
        self.conv4_0 = HarmonicResSEBlock(nb_filter[3], nb_filter[4], dilation=3)

        self.context = HybridContextModule(nb_filter[4])

        # Decoder dense skip connections with dual-domain fusion
        self.conv0_1 = FusionHarmonicBlock([nb_filter[0], nb_filter[1]], nb_filter[0])
        self.conv1_1 = FusionHarmonicBlock([nb_filter[1], nb_filter[2]], nb_filter[1])
        self.conv2_1 = FusionHarmonicBlock([nb_filter[2], nb_filter[3]], nb_filter[2])
        self.conv3_1 = FusionHarmonicBlock([nb_filter[3], nb_filter[4]], nb_filter[3])

        self.conv0_2 = FusionHarmonicBlock([nb_filter[0], nb_filter[0], nb_filter[1]], nb_filter[0])
        self.conv1_2 = FusionHarmonicBlock([nb_filter[1], nb_filter[1], nb_filter[2]], nb_filter[1])
        self.conv2_2 = FusionHarmonicBlock([nb_filter[2], nb_filter[2], nb_filter[3]], nb_filter[2])

        self.conv0_3 = FusionHarmonicBlock([nb_filter[0], nb_filter[0], nb_filter[0], nb_filter[1]], nb_filter[0])
        self.conv1_3 = FusionHarmonicBlock([nb_filter[1], nb_filter[1], nb_filter[1], nb_filter[2]], nb_filter[1])

        self.conv0_4 = FusionHarmonicBlock([nb_filter[0], nb_filter[0], nb_filter[0], nb_filter[0], nb_filter[1]], nb_filter[0])

        self.dropout = nn.Dropout2d(p=0.15)

        if self.deep_supervision:
            self.final1 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final2 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final3 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
            self.final4 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
        else:
            self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)

    def forward(self, input):
        # Encoder
        x0_0 = self.conv0_0(input)
        x1_0 = self.conv1_0(self.pool(x0_0))
        x2_0 = self.conv2_0(self.pool(x1_0))
        x3_0 = self.conv3_0(self.pool(x2_0))
        x4_0 = self.conv4_0(self.pool(x3_0))
        x4_0 = self.context(x4_0)

        # Decoder with gated dual-domain fusion
        x0_1 = self.conv0_1([x0_0, self.up(x1_0)])
        x1_1 = self.conv1_1([x1_0, self.up(x2_0)])
        x2_1 = self.conv2_1([x2_0, self.up(x3_0)])
        x3_1 = self.conv3_1([x3_0, self.up(x4_0)])

        x0_2 = self.conv0_2([x0_0, x0_1, self.up(x1_1)])
        x1_2 = self.conv1_2([x1_0, x1_1, self.up(x2_1)])
        x2_2 = self.conv2_2([x2_0, x2_1, self.up(x3_1)])

        x0_3 = self.conv0_3([x0_0, x0_1, x0_2, self.up(x1_2)])
        x1_3 = self.conv1_3([x1_0, x1_1, x1_2, self.up(x2_2)])

        x0_3 = self.dropout(x0_3)
        x1_3 = self.dropout(x1_3)

        x0_4 = self.conv0_4([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)])

        if self.deep_supervision:
            output1 = self.final1(x0_1)
            output2 = self.final2(x0_2)
            output3 = self.final3(x0_3)
            output4 = self.final4(x0_4)
            return [output1, output2, output3, output4]
        else:
            output = self.final(x0_4)
            return output
